package ants_pool

import (
	"context"
	"fmt"
	"github.com/panjf2000/ants/v2"
	"golang.org/x/sync/errgroup"
	"log"
	"sync"
)

// 封装的结构体
type PoolManager struct {
	pool   *ants.Pool
	g      *errgroup.Group
	ctx    context.Context
	cancel context.CancelFunc
	wg     sync.WaitGroup
}

// 创建新的池管理器
func NewPoolManager(poolSize int, maxTasks int, ctx0 context.Context) (*PoolManager, error) {
	pool, err := ants.NewPool(poolSize, ants.WithMaxBlockingTasks(maxTasks))
	if err != nil {
		return nil, err
	}

	ctx, cancel := context.WithCancel(ctx0)
	g, ctx := errgroup.WithContext(ctx)

	return &PoolManager{
		pool:   pool,
		g:      g,
		ctx:    ctx,
		cancel: cancel,
	}, nil
}

// 提交任务
func (pm *PoolManager) SubmitTask(task func() error) {
	pm.wg.Add(1)

	pm.g.Go(func() error {
		defer pm.wg.Done()

		// 使用 ants 池执行任务
		err := pm.pool.Submit(func() {
			if err := task(); err != nil {
				// 错误处理逻辑
				log.Printf("任务执行错误: %v", err)
			}
		})

		if err != nil {
			return fmt.Errorf("提交任务到协程池失败: %w", err)
		}
		return nil
	})
}

// 等待所有任务完成
func (pm *PoolManager) Wait() error {
	// 等待所有任务提交完成
	if err := pm.g.Wait(); err != nil {
		pm.cancel()
		return err
	}

	// 等待所有任务执行完成
	pm.wg.Wait()
	return nil
}

// 关闭池
func (pm *PoolManager) Close() {
	pm.cancel()
	pm.pool.Release()
}
